import imaplib
import email

from saver import *

version = sys.version_info.major


def login(host, username, password, port=993):
    connection = imaplib.IMAP4_SSL(host, port)
    connection.login(username, password)
    return connection


def get_mail_list(connection):
    connection.select('INBOX')
    result, data = connection.search(None, 'ALL')
    return data[0].split()


def get_attachment(connection, mail_list):
    for mail in mail_list:
        result, data = connection.fetch(mail, '(RFC822)')
        if version == 2:
            msg = email.message_from_string(data[0][1])
        else:
            msg = email.message_from_string(data[0][1].decode())
        for part in msg.walk():
            if not part.is_multipart():
                name = part.get_filename()  # incorrect decode
                charset = part.get_charset()
                if name:
                    if version == 2:
                        header = email.Header.Header(name)
                        header = email.Header.decode_header(header)
                    else:
                        header = email.header.Header(name)
                        header = email.header.decode_header(header)
                    filename = header[0][0]
                    filename = filename.decode(header[0][1], 'gbk')
                    # fixme: filename should support unicode
                    print("get attachment: %s" % filename)
                    filedata = part.get_payload(decode=True)
                    # filename = bytes.decode(filename)
                    save_attachment(filename, filedata)
                else:
                    data = part.get_payload(decode=True)


